Skip to content

Add rampup batch size support in MaxText #2535

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 1 commit into from
Closed

Add rampup batch size support in MaxText #2535

wants to merge 1 commit into from

Conversation

@NuojCheng
Copy link
Collaborator

@NuojCheng NuojCheng commented Oct 22, 2025 •
edited
Loading

Description

This PR adds support for ramp-up batch size, a feature originally proposed in the GPT-3 paper and implemented in Megatron.

When enabled, the per device batch size starts at a smaller value (per_device_batch_size_start) and gradually increases (per_device_batch_size_increment) until it reaches the target per_device_batch_size over a specified number of rampup_samples. This can help improve training stability, especially during the initial training phases.

This feature introduces four new configuration parameters, which align with the Megatron implementation:

  • enable_rampup_batch_size: (default: False) Set to True to enable the ramp-up feature.
  • per_device_batch_size_start: The per-device batch size to use at the beginning of training.
  • per_device_batch_size_increment: The amount to increase the per-device batch size at each ramp-up step.
  • global_rampup_samples: The total number of samples to process before reaching the full target batch size.

The PR includes the following changes:

  • RampupDataLoader: Adds a new RampupDataLoader class that inherits from the base DataLoader. Its primary responsibility is to truncate the input data to match the correct ramp-up shape for the current training step.
  • Metric Logger: Updates the metric logger to prevent flops and token counts associated with metadata from being logged.
  • Config Updates: Modifies pyconfig.py to register and validate the new ramp-up configuration parameters.
  • Testing: Adds new tests to data_loader_tests.py to verify the RampupDataLoader's slicing and increment logic.

FIXES: b/452468482

Tests

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@NuojCheng NuojCheng added the draft Draft PR label Oct 22, 2025
@NuojCheng NuojCheng changed the title Add rampup batch size support in MaxText [WIP] Add rampup batch size support in MaxText Oct 22, 2025
gobbleturk
gobbleturk reviewed Oct 23, 2025
View reviewed changes
@NuojCheng NuojCheng force-pushed the chengnuojin-rampup-batch branch 4 times, most recently from 57ff3e8 to 842193d Compare October 24, 2025 00:53
@NuojCheng NuojCheng changed the title [WIP] Add rampup batch size support in MaxText Add rampup batch size support in MaxText Oct 24, 2025
@NuojCheng NuojCheng added gemini-review and removed draft Draft PR labels Oct 24, 2025
@github-actions
Copy link

github-actions bot commented Oct 24, 2025

🤖 Hi @NuojCheng, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

@NuojCheng NuojCheng force-pushed the chengnuojin-rampup-batch branch from 842193d to 82f5edf Compare October 24, 2025 20:48
@NuojCheng NuojCheng marked this pull request as ready for review October 24, 2025 20:49
@RissyRan
Copy link
Collaborator

RissyRan commented Oct 29, 2025

🤖 Hi @NuojCheng, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

It seems out of quota for free tier. We are going to update the Tier 1, should be better soon.

Attempt 1 failed with status 429. Retrying with backoff... ApiError: {"error":{"message":"{\n  \"error\": {\n    \"code\": 429,\n    \"message\": \"You exceeded your current quota, please check your plan and billing details. For more information on this error, head to: https://ai.google.dev/gemini-api/docs/rate-limits. To monitor your current usage, head to: https://ai.dev/usage?tab=rate-limit.\\n* Quota exceeded for metric: generativelanguage.googleapis.com/generate_content_free_tier_requests, limit: 2\\nPlease retry in 20.529613201s.\",\n    \"status\": \"RESOURCE_EXHAUSTED\",\n    \"details\": [\n      {\n        \"@type\": \"type.googleapis.com/google.rpc.QuotaFailure\",\n        \"violations\": [\n          {\n            \"quotaMetric\": \"generativelanguage.googleapis.com/generate_content_free_tier_requests\",\n            \"quotaId\": \"GenerateRequestsPerMinutePerProjectPerModel-FreeTier\",\n            \"quotaDimensions\": {\n              \"location\": \"global\",\n              \"model\": \"gemini-2.5-pro\"\n            },\n            \"quotaValue\": \"2\"\n          }\n        ]\n      },\n      {\n        \"@type\": \"type.googleapis.com/google.rpc.Help\",\n        \"links\": [\n          {\n            \"description\": \"Learn more about Gemini API quotas\",\n            \"url\": \"[https://ai.google.dev/gemini-api/docs/rate-limits\](https://ai.google.dev/gemini-api/docs/rate-limits/)"\n          }\n        ]\n      },\n      {\n        \"@type\": \"type.googleapis.com/google.rpc.RetryInfo\",\n        \"retryDelay\": \"20s\"\n      }\n    ]\n  }\n}\n","code":429,"status":"Too Many Requests"}}

@github-actions
Copy link

github-actions bot commented Oct 29, 2025

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

@NuojCheng NuojCheng force-pushed the chengnuojin-rampup-batch branch 2 times, most recently from 4114fa7 to b911ca8 Compare October 31, 2025 02:14
gobbleturk
gobbleturk reviewed Oct 31, 2025
View reviewed changes
@github-actions
Copy link

github-actions bot commented Oct 31, 2025

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

github-actions[bot]
github-actions bot reviewed Oct 31, 2025
View reviewed changes
Copy link

@github-actions github-actions bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📋 Review Summary

This pull request introduces a batch size ramp-up feature to improve training stability, which is a valuable addition. The implementation is generally clean and follows existing patterns in the codebase. The configuration, data loading, and metric logging components are well-integrated.

🔍 General Feedback

  • The refactoring of the sharding logic out of the DataLoader and into the training loop is a good improvement for separation of concerns.
  • The use of a factory function create_dataloader is a clean way to handle the conditional creation of the RampUpDataLoader.
  • I've left a couple of minor suggestions to improve an assertion message and to correct the logic in a unit test.

RissyRan
RissyRan reviewed Oct 31, 2025
View reviewed changes
Copy link
Collaborator

@RissyRan RissyRan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks Nuojing! LGTM! If you could response to Gemini's comments, that will be great! We could chat more on future deliveries offline, but not a blocker for this PR.

@NuojCheng NuojCheng force-pushed the chengnuojin-rampup-batch branch 3 times, most recently from 24e0d29 to 5e7f164 Compare October 31, 2025 20:07
RissyRan
RissyRan approved these changes Oct 31, 2025
View reviewed changes
Copy link
Collaborator

@RissyRan RissyRan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM for this PR! let's chat more on remaining tasks in the meeting today. Thank you!

aireenmei
aireenmei approved these changes Oct 31, 2025
View reviewed changes
@NuojCheng NuojCheng force-pushed the chengnuojin-rampup-batch branch 2 times, most recently from 47a1028 to d3fed21 Compare November 1, 2025 00:09
@NuojCheng NuojCheng force-pushed the chengnuojin-rampup-batch branch from d3fed21 to 40e056d Compare November 1, 2025 00:12
@RissyRan
Copy link
Collaborator

RissyRan commented Nov 2, 2025

We discussed offline, and agree to merge this feature request first. Leave 2 action items for follow up: 1) learning rate adjustment if needed due to this feature; 2) metrics report like throughput during the batch ramp up.

@NuojCheng NuojCheng force-pushed the chengnuojin-rampup-batch branch 3 times, most recently from 3375760 to 9143004 Compare November 2, 2025 06:02
copybara-service bot pushed a commit that referenced this pull request Nov 2, 2025
@NuojCheng @Google-ML-Automation
Copybara import of the project:
d2f608d 
--
40e056d by NuojCheng <[email protected]>:

add rampup batch size

COPYBARA_INTEGRATE_REVIEW=#2535 from AI-Hypercomputer:chengnuojin-rampup-batch 40e056d
PiperOrigin-RevId: 827037473
gulsumgudukbay pushed a commit to ROCm/maxtext that referenced this pull request Nov 8, 2025
@NuojCheng @gulsumgudukbay
Copybara import of the project:
fcd1850 
--
40e056d by NuojCheng <[email protected]>:

add rampup batch size

COPYBARA_INTEGRATE_REVIEW=AI-Hypercomputer#2535 from AI-Hypercomputer:chengnuojin-rampup-batch 40e056d
PiperOrigin-RevId: 827037473
gulsumgudukbay pushed a commit to ROCm/maxtext that referenced this pull request Nov 10, 2025
@NuojCheng @gulsumgudukbay
Copybara import of the project:
9ea94bc 
--
40e056d by NuojCheng <[email protected]>:

add rampup batch size

COPYBARA_INTEGRATE_REVIEW=AI-Hypercomputer#2535 from AI-Hypercomputer:chengnuojin-rampup-batch 40e056d
PiperOrigin-RevId: 827037473
@NuojCheng NuojCheng force-pushed the chengnuojin-rampup-batch branch from 9143004 to 3da3587 Compare November 11, 2025 01:45
@NuojCheng NuojCheng closed this Nov 11, 2025
melissawm pushed a commit to melissawm/maxtext that referenced this pull request Nov 18, 2025
@NuojCheng @melissawm
Copybara import of the project:
f4a2b0f 
--
40e056d by NuojCheng <[email protected]>:

add rampup batch size

COPYBARA_INTEGRATE_REVIEW=AI-Hypercomputer#2535 from AI-Hypercomputer:chengnuojin-rampup-batch 40e056d
PiperOrigin-RevId: 827037473
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@github-actions github-actions[bot] github-actions[bot] left review comments

@gobbleturk gobbleturk gobbleturk left review comments

@aireenmei aireenmei aireenmei approved these changes

@RissyRan RissyRan RissyRan approved these changes

@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner

@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner

@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner

@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner

@gagika gagika Awaiting requested review from gagika gagika is a code owner

@shralex shralex Awaiting requested review from shralex shralex is a code owner

@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC is a code owner

@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner

@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner

@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner

@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner

@NicoGrande NicoGrande Awaiting requested review from NicoGrande NicoGrande is a code owner

Assignees

No one assigned

Labels

pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@NuojCheng @RissyRan @aireenmei @gobbleturk